import os

def remove_useless_labels(txt_folder, remove_labels):
    # 获取txt文件夹中的所有文件
    txt_files = [f for f in os.listdir(txt_folder) if f.endswith('.txt')]

    # 遍历每个txt文件
    for txt_file in txt_files:
        # 读取txt文件内容
        with open(os.path.join(txt_folder, txt_file), 'r') as f:
            lines = f.readlines()

        # 删除不需要的标签
        modified_lines = []
        for line in lines:
            parts = line.split()
            label = parts[0]

            # 删除不需要的标签
            if label not in remove_labels:
                modified_lines.append(line)

        # 将修改后的内容写入txt文件
        with open(os.path.join(txt_folder, txt_file), 'w') as f:
            f.writelines(modified_lines)

# 示例用法
useless_labels = ['2', '6','7','9']  # 需要删除的标签列表
txt_folder = 'D:\\datasets\\visDrone\\VisDrone2019-DET-val\\new_label_0.01'  # txt文件夹路径


# 执行删除操作
remove_useless_labels(txt_folder, useless_labels)